#include "Instance.h"
#include "Material/GGXMaterial.h"
#include "Material/LambertianMaterial.h"
#include "Material/MixedMaterial.h"
#include "Mesh/TriangleMesh.h"
#include "core/Scene.h"

using namespace OBR;

bool __tryLoadTextureMap(std::string path, ColorMipmap &dest)
{
    bool success = false;
    try
    {
        Json::Value cfg;
        cfg["type"] = "image";
        cfg["path"] = path;
        ColorTexture2D tex;
        ConfigHelpers::loadTexture2D<Spectrum>(cfg, tex);
        dest.createFromTexture(tex);
        success = true;
    }
    catch (const std::exception &e)
    {
        std::cerr << "[error loading texture map]: " << e.what() << '\n';
    }
    return success;
}

// TODO: move to helpers
bool __loadOBJFileAsStaticInstance(std::string path, StaticInstanceBuildData &result)
{
    // check path
    if (path.length() == 0)
    {
        std::cerr << "path provided to OBJLoader cannot be empty" + path << std::endl;
        return false;
    }

    bool success = false;
    objl::Loader loader;
    // load file
    success = loader.LoadFile(path);
    if (success)
    {
        std::string obj_dir = path.substr(0, path.find_last_of("/") + 1);
        for (objl::Mesh &m : loader.LoadedMeshes)
        {
            std::shared_ptr<TriangleMesh> mesh = std::make_shared<TriangleMesh>();
            // vertex
            for (unsigned int &idx : m.Indices)
            {
                objl::Vertex v = m.Vertices[idx];
                Vertex vert;
                vert.n = make_float3(v.Normal.X, v.Normal.Y, v.Normal.Z);
                vert.pos = make_float3(v.Position.X, v.Position.Y, v.Position.Z);
                vert.tex_uv = make_float2(v.TextureCoordinate.X, v.TextureCoordinate.Y);
                mesh->vertexes.push_back(vert);
            }
            // triangle
            for (int base_offset = 0; base_offset < m.Indices.size(); base_offset += 3)
            {
                Triangle tri;
                tri.vertexes[0] = m.Indices[base_offset];
                tri.vertexes[1] = m.Indices[base_offset + 1];
                tri.vertexes[2] = m.Indices[base_offset + 2];
                mesh->triangles.push_back(tri);
            }
            mesh->updateAABB();
            // material
            objl::Material rawmat = m.MeshMaterial;
            Json::Value tex_cfg;
            tex_cfg["type"] = "image";
            // Kd part using lambert material
            std::shared_ptr<LambertianMaterial> kd_mat = std::make_shared<LambertianMaterial>();
            if (rawmat.map_Kd.length() == 0 ||
                !__tryLoadTextureMap(obj_dir + rawmat.map_Kd, kd_mat->albedo))
            {
                // single color
                Spectrum color = Spectrum(0.0f);
                color.fromRGB(rawmat.Kd.X * 256, rawmat.Kd.Y * 256, rawmat.Kd.Z * 256);
                kd_mat->albedo.createFromTexture(color);
            }
            // Ks part using GGX material
            // @https://computergraphics.stackexchange.com/questions/1515/what-is-the-accepted-method-of-converting-shininess-to-roughness-and-vice-versa
            std::shared_ptr<GGXMaterial> ks_mat = std::make_shared<GGXMaterial>();
            float roughness = powf(2.0f / (2.0f + rawmat.Ns), 0.5);
            ks_mat->alpha.createFromTexture(GGXMaterial::roughnessToAlpha(roughness));
            ks_mat->eta = 1.5f;
            if (rawmat.map_Ks.length() == 0 ||
                !__tryLoadTextureMap(obj_dir + rawmat.map_Ks, ks_mat->albedo))
            {
                // single color
                Spectrum color = Spectrum(0.0f);
                color.fromRGB(rawmat.Ks.X * 256, rawmat.Ks.Y * 256, rawmat.Ks.Z * 256);
                ks_mat->albedo.createFromTexture(color);
            }
            // combine material
            std::shared_ptr<MixedMaterial> mat = std::make_shared<MixedMaterial>();
            mat->mat_list.push_back(kd_mat);
            mat->mat_list.push_back(ks_mat);
            mat->weights.push_back(1.0f);
            mat->weights.push_back(1.0f);

            result.materials.push_back(kd_mat);
            result.materials.push_back(ks_mat);
            result.materials.push_back(mat);
            result.meshes.push_back(mesh);
        }
    }
    else
    {
        std::cerr << "error loading obj file @ " + path << std::endl;
    }
    return success;
}

Instance::Instance() : ASNode(ASNodeType::IAS)
{
}

Instance::~Instance()
{
}

/* ------------------------ static member instancing ------------------------ */
std::vector<HitData> StaticInstance::hit_data;
std::vector<std::shared_ptr<Mesh>> StaticInstance::hit_meshes;
std::vector<size_t> StaticInstance::hit_gas_offset;

StaticInstance::StaticInstance()
{
}

StaticInstance::~StaticInstance()
{
}

void StaticInstance::linkChild(VirtualProgramManager materials, VirtualProgramManager mediums,
                               VirtualProgramManager meshes, VirtualProgramManager lights,
                               std::vector<std::shared_ptr<Instance>> instances)
{
    for (uint childid = 0; childid < this->child_data.size(); childid++)
    {
        const InstanceChildData &cdata = this->child_data[childid];

        std::shared_ptr<ASNode> child = nullptr;
        std::shared_ptr<AreaLight> light = nullptr;
        if (cdata.type == ASNodeType::GAS)
        {
            child = std::dynamic_pointer_cast<ASNode>(meshes[cdata.asid]);
        }
        else
        {
            child = instances[cdata.asid];
        }

        // generate hit id
        size_t hit_id = -1;
        if (child->type == ASNodeType::GAS)
        {
            hit_id = StaticInstance::hit_data.size();
            // create hit data
            std::shared_ptr<Mesh> mesh = std::dynamic_pointer_cast<Mesh>(child);
            HitData data;
            data.mesh_id = mesh->program_id;
            data.mesh_instance_id = mesh->program_instance_id;
            if (cdata.matid != -1)
            {
                data.mat_id = materials[cdata.matid]->program_id;
                data.mat_instance_id = materials[cdata.matid]->program_instance_id;
            }
            else
            {
                data.mat_id = -1;
            }
            if (cdata.lightid != -1)
            {
                light = std::dynamic_pointer_cast<AreaLight>(lights[cdata.lightid]);
                data.light_id = light->program_id;
                data.light_instance_id = light->program_instance_id;
                light->linkMesh(mesh);
            }
            else
            {
                data.light_id = -1;
            }
            if (cdata.medid[0] != -1)
            {
                data.med_id_inside = mediums[cdata.medid[0]]->program_id;
                data.med_instance_id_inside = mediums[cdata.medid[0]]->program_instance_id;
            }
            else
            {
                data.med_id_inside = -1;
            }
            if (cdata.medid[1] != -1)
            {
                data.med_id_outside = mediums[cdata.medid[1]]->program_id;
                data.med_instance_id_outside = mediums[cdata.medid[1]]->program_instance_id;
            }
            else
            {
                data.med_id_outside = -1;
            }
            // add to global array
            StaticInstance::hit_data.push_back(data);
            StaticInstance::hit_meshes.push_back(mesh);
            StaticInstance::hit_gas_offset.push_back(child->ref_cnt);
        }
        this->hit_ids.push_back(hit_id);
        this->surface_lights.push_back(light);

        // inc refcount of child and push child to array
        child->ref_cnt++;
        this->children.push_back(child);

        // enlarge aabb
        this->aabb = child == 0 ? child->aabb : Bounds::combine(this->aabb, child->aabb);
    }
}

void StaticInstance::buildAccelStructure(OptixDeviceContext &ctx, CUstream &stream)
{
    if (this->built)
    {
        return;
    }

    /* --------------------------- build children --------------------------- */
    for (auto &&c : this->children)
    {
        c->buildAccelStructure(ctx, stream);
        if (!c->built)
        {
            std::cerr << "children cannot be built" << std::endl;
            throw std::runtime_error("children cannot be built");
        }
    }

    const size_t n_instances = this->children.size();
    const size_t instances_size = sizeof(OptixInstance) * n_instances;

    /* --------------------- create instances and upload -------------------- */
    OptixInstance *instances = new OptixInstance[n_instances];
    for (size_t i = 0; i < n_instances; i++)
    {
        instances[i].flags = OPTIX_INSTANCE_FLAG_DISABLE_ANYHIT;
        instances[i].visibilityMask = 1;
        instances[i].instanceId = i;
        instances[i].traversableHandle = this->children[i]->handle;
        if (this->children[i]->type == ASNodeType::GAS)
        {
            instances[i].sbtOffset = StaticInstance::hit_meshes[this->hit_ids[i]]->gas_sbt_base +
                                     StaticInstance::hit_gas_offset[this->hit_ids[i]];
        }
        else
        {
            instances[i].sbtOffset = 0;
        }
        memcpy(instances[i].transform, this->transforms[i].data, sizeof(float) * 12);
    }
    this->d_instances.resize(instances_size);
    this->d_instances.upload(instances, instances_size);
    delete[] instances;

    /* ----------------------------- build accel ---------------------------- */

    OptixBuildInput instance_input = {};
    instance_input.type = OPTIX_BUILD_INPUT_TYPE_INSTANCES;
    instance_input.instanceArray.instances = d_instances.d_pointer();
    instance_input.instanceArray.numInstances = n_instances;

    OptixAccelBuildOptions accel_options = {};
    accel_options.buildFlags = OPTIX_BUILD_FLAG_NONE;
    accel_options.operation = OPTIX_BUILD_OPERATION_BUILD;

    accel_options.motionOptions.numKeys = 1;
    accel_options.motionOptions.timeBegin = 0.0f;
    accel_options.motionOptions.timeEnd = 1.0f;
    accel_options.motionOptions.flags = OPTIX_MOTION_FLAG_NONE;

    OptixAccelBufferSizes ias_buffer_sizes;
    OPTIX_CHECK(optixAccelComputeMemoryUsage(ctx, &accel_options, &instance_input,
                                             1, // num build inputs
                                             &ias_buffer_sizes));

    CudaBuffer d_temp_buffer(ias_buffer_sizes.tempSizeInBytes);
    this->ias_output_buffer.resize(ias_buffer_sizes.outputSizeInBytes);

    OPTIX_CHECK(optixAccelBuild(ctx,
                                stream, // CUDA stream
                                &accel_options, &instance_input,
                                1, // num build inputs
                                d_temp_buffer.d_pointer(), ias_buffer_sizes.tempSizeInBytes,
                                this->ias_output_buffer.d_pointer(),
                                ias_buffer_sizes.outputSizeInBytes, &this->handle,
                                nullptr, // emitted property list
                                0        // num emitted properties
                                ));

    this->built = true;
}

void StaticInstance::_buildFromConfig(const Json::Value &config)
{
    StaticInstanceBuildData data;
    if (config.isMember("obj source"))
    {
        std::string source_path = config["obj source"].asString();
        if (source_path.rfind("./", 0) == 0)
        {
            source_path = std::string(PROJ_ROOT) + source_path.substr(1, source_path.size() - 1);
        }
        if (!__loadOBJFileAsStaticInstance(source_path, data))
        {
            std::cerr << "failed to load obj file" << std::endl;
            throw std::runtime_error("failed to load obj file");
        }
        data.medid = config.isMember("Medium") ? config["Medium"].asInt() : -1;
    }
    else
    {
        for (Json::ArrayIndex i = 0; i < config["children"].size(); i++)
        {
            Json::Value instance_config = config["children"][i];
            InstanceChildData child;

            /* ---------------------- set child properties ---------------------- */
            if (instance_config["type"].asString().compare("gas") == 0)
            {
                child.type = ASNodeType::GAS;
                child.asid = instance_config["mesh id"].asUInt();
            }
            else if (instance_config["type"].asString().compare("ias") == 0)
            {
                child.type = ASNodeType::IAS;
                child.asid = instance_config["instance id"].asUInt();
            }
            else
            {
                std::cerr << "invalid child config type" << std::endl;
                throw std::runtime_error("invalid child config type");
            }
            Matrix4x4f transform = ConfigHelpers::loadTransformMatrix(instance_config["transform"]);

            /* --------------------- set optional properties -------------------- */
            if (instance_config.isMember("material"))
            {
                child.matid = instance_config["material"].asInt();
            }
            if (instance_config.isMember("light"))
            {
                child.lightid = instance_config["light"].asInt();
            }
            if (instance_config.isMember("medium inside"))
            {
                child.medid[0] = instance_config["medium inside"].asInt();
            }
            if (instance_config.isMember("medium outside"))
            {
                child.medid[1] = instance_config["medium outside"].asInt();
            }

            // add to list
            this->child_data.push_back(child);
            this->transforms.push_back(transform);
        }
    }
    this->build_data = data;
}

void StaticInstance::postBuild(Scene &sc)
{
    for (auto &&mat : this->build_data.materials)
    {
        sc.materials.addInstance(mat);
    }
    for (int i = 0; i < this->build_data.meshes.size(); i++)
    {
        std::shared_ptr<Mesh> mesh = this->build_data.meshes[i];
        std::shared_ptr<Material> mat = this->build_data.materials[i * 3 + 2];
        // add to scene mesh
        sc.meshes.addInstance(mesh);
        // link mesh child
        HitData hdata;
        hdata.mesh_id = mesh->program_id;
        hdata.mesh_instance_id = mesh->program_instance_id;
        hdata.mat_id = mat->program_id;
        hdata.mat_instance_id = mat->program_instance_id;
        if (this->build_data.medid != -1)
        {
            hdata.med_id_inside = sc.mediums[this->build_data.medid]->program_id;
            hdata.med_instance_id_inside = sc.mediums[this->build_data.medid]->program_instance_id;
            hdata.med_id_outside = sc.mediums[this->build_data.medid]->program_id;
            hdata.med_instance_id_outside = sc.mediums[this->build_data.medid]->program_instance_id;
        }
        // to hit group
        uint hit_id = StaticInstance::hit_data.size();
        StaticInstance::hit_data.push_back(hdata);
        StaticInstance::hit_meshes.push_back(mesh);
        StaticInstance::hit_gas_offset.push_back(mesh->ref_cnt);

        mesh->ref_cnt++;
        this->hit_ids.push_back(hit_id);
        this->children.push_back(mesh);
        this->surface_lights.push_back(nullptr);
        this->transforms.push_back(Matrix4x4f(1.0f));
        this->aabb = i == 0 ? mesh->aabb : Bounds::combine(this->aabb, mesh->aabb);
    }

    if (this->build_data.meshes.size() == 0)
    {
        this->linkChild(sc.materials, sc.mediums, sc.meshes, sc.lights, sc.instances);
    }
}

void StaticInstance::updateLocalToWorld(Matrix4x4f l2w)
{
    uint n_children = this->children.size();
    for (uint i = 0; i < n_children; i++)
    {
        Matrix4x4f current = Matrix4x4f::matmult(l2w, this->transforms[i]);
        if (this->children[i]->type != ASNodeType::GAS)
        {
            this->children[i]->updateLocalToWorld(current);
        }
        else if (this->surface_lights[i] != nullptr)
        {
            this->surface_lights[i]->local2world = current;
        }
    }
}
